import shutil

import numpy
import numpy as np
import cv2
import os
from torch.utils.data import Dataset, DataLoader
import torch
from utils import draw_gaussian, getResizeCoordinate, getResizePos, scale_value
import sys
from skimage import exposure
import pandas as pd
from utils import calc_vector


def createHeatMap(landmarks, heatmap_size):
    temp_landmarks = landmarks.reshape(-1, 2)
    heatmap = np.zeros((temp_landmarks.shape[0], heatmap_size, heatmap_size), dtype=np.float32)
    for i in range(temp_landmarks.shape[0]):
        heatmap[i] = draw_gaussian(heatmap[i], temp_landmarks[i]+1, 2)
    return heatmap


def gazeto2d(gaze):
    yaw = np.arctan2(-gaze[0], -gaze[2])
    pitch = np.arcsin(-gaze[1])
    return np.array([yaw, pitch])


class loader(Dataset):
    def __init__(self, path, root, header=True):
        self.lines = []
        # if isinstance(path, list):
        #     for i in path:
        #         with open(i) as f:
        #             line = f.readlines()
        #             if header: line.pop(0)
        #             self.lines.extend(line)
        # else:
        #     with open(path) as f:
        #         self.lines = f.readlines()
        #         if header: self.lines.pop(0)

        if isinstance(path, list):
            for label in path:
                line1 = np.asarray(pd.read_csv(label, header=None))
                if header: line1.pop(0)
                self.lines.extend(line1)
        else:
            self.lines = np.asarray(pd.read_csv(path, header=None))
        self.root = root
        # self.class_num = int(sys.argv[3])

    def __len__(self):
        return len(self.lines)

    def __getitem__(self, idx):

        line = self.lines[idx]
        # line = line.strip().split(" ")

        gaze2d = line[5]
        # is_train = sys.argv[0].split("/")[-1] == "train.py"
        # face_name = line[0]
        # point_p = file_name.split("\\")[-1].replace("jpg", "csv")
        # point_path = "/home/xian/Downloads/Gaze-Net-main/Dataset/Gaze_360_p_eye/train/" + point_p
        # point_csv = np.asarray(pd.read_csv(point_path))
        # face_point_path = "/home/xian/Downloads/Gaze-Net-main/Dataset/gaze360_train_face/" + point_p
        # point_csv_face = np.asarray(pd.read_csv(face_point_path))

        landmarks = []
        # left
        # t_x = 0
        # t_y = 0
        # for point_i in range(10, 10 + 8):
        #     t_x += point_csv[0][point_i]
        #     t_y += point_csv[0][point_i + 56]
        # landmarks.append(np.asarray([t_x / 8, t_y / 8]))

        # for point_i in range(10+8, 10 + 20):
        #     t_x = point_csv[0][point_i]
        #     t_y = point_csv[0][point_i + 56]
        #     landmarks.append(np.asarray([t_x, t_y]))

        # right
        # t_x = 0
        # t_y = 0
        # for point_i in range(10 + 28, 10 + 36):
        #     t_x += point_csv[0][point_i]
        #     t_y += point_csv[0][point_i + 56]
        # landmarks.append(np.asarray([t_x / 8, t_y / 8]))

        # for point_i in range(10 + 36, 10 + 48):
        #     t_x = point_csv[0][point_i]
        #     t_y = point_csv[0][point_i + 56]
        #     landmarks.append(np.asarray([t_x, t_y]))

        # landmarks.append(np.asarray([point_csv_face[0][2 + 36], point_csv_face[0][2 + 36 + 68]]))
        # landmarks.append(np.asarray([point_csv_face[0][2 + 39], point_csv_face[0][2 + 39 + 68]]))
        # landmarks.append(np.asarray([point_csv_face[0][2 + 42], point_csv_face[0][2 + 42 + 68]]))
        # landmarks.append(np.asarray([point_csv_face[0][2 + 45], point_csv_face[0][2 + 45 + 68]]))
        # 
        # landmarks.append(np.asarray([point_csv_face[0][2 + 27], point_csv_face[0][2 + 27 + 68]]))
        # landmarks.append(np.asarray([point_csv_face[0][2 + 28], point_csv_face[0][2 + 28 + 68]]))
        # landmarks.append(np.asarray([point_csv_face[0][2 + 29], point_csv_face[0][2 + 29 + 68]]))
        # landmarks.append(np.asarray([point_csv_face[0][2 + 30], point_csv_face[0][2 + 30 + 68]]))

        # for point_i in range(10, 10 + 56):
        #     t_x = point_csv[0][point_i]
        #     t_y = point_csv[0][point_i + 56]
        #     landmarks.append(np.asarray([t_x, t_y]))
        # for point_i in range(68):
        #     landmarks.append(np.asarray([point_csv_face[0][2 + point_i], point_csv_face[0][2 + point_i + 68]]))
        landmarks.append(np.asarray([line[6 + 468 * 2], line[6 + 468 * 2 + 1]]).astype(float))
        landmarks.append(np.asarray([line[6 + 473 * 2], line[6 + 473 * 2 + 1]]).astype(float))
        landmarks.append(np.asarray([line[6 + 130 * 2], line[6 + 130 * 2 + 1]]).astype(float))  # 左眼左眼角
        landmarks.append(np.asarray([line[6 + 243 * 2], line[6 + 243 * 2 + 1]]).astype(float))  # 左眼右眼角
        landmarks.append(np.asarray([line[6 + 463 * 2], line[6 + 463 * 2 + 1]]).astype(float))  # 右眼左眼角
        landmarks.append(np.asarray([line[6 + 359 * 2], line[6 + 359 * 2 + 1]]).astype(float))  # 右眼右眼角
        landmarks.append(np.asarray([line[6 + 27 * 2], line[6 + 27 * 2 + 1]]).astype(float))  # 左眼上
        landmarks.append(np.asarray([line[6 + 23 * 2], line[6 + 23 * 2 + 1]]).astype(float))  # 左眼下
        landmarks.append(np.asarray([line[6 + 257 * 2], line[6 + 257 * 2 + 1]]).astype(float))  # 右眼上
        landmarks.append(np.asarray([line[6 + 253 * 2], line[6 + 253 * 2 + 1]]).astype(float))  # 右眼下
        # landmarks.append(np.asarray([line[6 + 159 * 2], line[6 + 159 * 2 + 1]]).astype(float))  # 左眼皮上
        # landmarks.append(np.asarray([line[6 + 386 * 2], line[6 + 386 * 2 + 1]]).astype(float))  # 右眼皮上

        landmarks = np.array(landmarks)

        new_landmarks = getResizePos(landmarks, [224, 224], [56, 56])
        heatmap = createHeatMap(new_landmarks, 56)

        face = line[0]

        label_t = np.array(gaze2d.split(",")).astype("float")
        label_t = np.rad2deg(label_t)# + 180.0) / 360.0
        gaze_rem = label_t
        gaze_rem_reg = gaze_rem

        # if is_train:
        #     gaze_rem = gaze_rem + 90
        #     gaze_rem = gaze_rem.round()
        #     if gaze_rem[0] >= 180:
        #         gaze_rem[0] = 179
        #     if gaze_rem[0] <= 0:
        #         gaze_rem[0] = 0
        #     if gaze_rem[1] >= 180:
        #         gaze_rem[1] = 179
        #     if gaze_rem[1] <= 0:
        #         gaze_rem[1] = 0

        img_dir = face.replace("\\", "/")
        img_dir = img_dir.split("/")
        new_path = '/'.join([img_dir[0], 'Face', img_dir[1]])
        # fimg = cv2.imread(os.path.join(self.root, face))
        fimg = cv2.imread(os.path.join(self.root, new_path)) # / 255.0

        # for cc_i, tem_ii in enumerate(landmarks):
        #     cv2.circle(fimg, (int(tem_ii[0]), int(tem_ii[1])), 1, (0, 255, 255), 1)
        #     cv2.imshow("test", fimg)
        #     cv2.waitKey(0)


        fimg = fimg.transpose(2, 0, 1)

        img = {
            "face": torch.from_numpy(fimg).type(torch.FloatTensor),
        }
        target = {
            "target_reg": torch.from_numpy(gaze_rem).type(torch.FloatTensor),
            "target_reg_reg": torch.from_numpy(gaze_rem_reg).type(torch.FloatTensor),
            "landmarks": torch.from_numpy(landmarks).type(torch.FloatTensor),
            "new_landmarks": torch.from_numpy(new_landmarks).type(torch.FloatTensor),
            "heatmap": torch.from_numpy(heatmap).type(torch.FloatTensor),
            "img_name": os.path.join(self.root, face),
        }

        return img, target


def txtload(labelpath, imagepath, batch_size, shuffle=True, num_workers=0, header=True):
    # labelpath = ["/home/xian/Downloads/Gaze-Net-main/Dataset/FaceBased/Gaze360/Label/test_eye_face.label",
    #              "/home/xian/Downloads/Gaze-Net-main/Dataset/FaceBased/Gaze360/Label/train_eye_face.label"]#["/home/ys/Downloads/gaze_eye_diap/Gaze-Net-main/Dataset/FaceBased/Gaze360/Label/train.label", "/home/ys/Downloads/gaze_eye_diap/Gaze-Net-main/Dataset/FaceBased/Gaze360/Label/test.label"]
    imagepath = "/home/xian/mzs/mzs_code/Dataset/360/Image"
    labelpath = "/home/xian/mzs/mzs_code/Dataset/360/label/middle_360.csv"
    dataset = loader(labelpath, imagepath, header)
    print(f"[Read Data]: Total num: {len(dataset)}")
    print(f"[Read Data]: Label path: {labelpath}")
    load = DataLoader(dataset, batch_size=batch_size, shuffle=shuffle, num_workers=num_workers)
    return load

if __name__ == "__main__":
    path = '/home/ys/Downloads/Gaze-Net-main/Dataset/FaceBased/Gaze360/Label/p5000_del.label'
    d = loader(path, "/home/ys/Downloads/Gaze-Net-main/Dataset/FaceBased/Gaze360/Image")
    print(len(d))
    for i in range(len(d)):
        (data, label) = d.__getitem__(i)
        target_reg = label['target_reg']
        yaw = target_reg[0]
        # pitch = target_reg[1]

        gaze_rem = yaw + 90
        gaze_rem = gaze_rem.round()
        if gaze_rem >= 180:
            gaze_rem = 179
        if gaze_rem <= 0:
            gaze_rem = 0
        #print(int(np.asarray(gaze_rem)))
        temp_class = int(np.asarray(gaze_rem))
        path = str(temp_class)
        file_old = label['img_name']
        file_old = file_old.replace("Face/", "Right/")
        t = file_old.split("/")
        file_name = t[-3] + "_" + t[-2] + "_" + t[-1]
        file_new = os.path.join(path, file_name)
        if os.path.exists(path):
            shutil.copyfile(file_old, file_new)
        else:
            os.mkdir(path)
            shutil.copyfile(file_old, file_new)


    # path = '/home/ys/Downloads/Gaze-Net-main/Dataset/FaceBased/Gaze360/Label/p5000_del.label'
    # d = loader(path, "/home/ys/Downloads/Gaze-Net-main/Dataset/FaceBased/Gaze360/Image")
    # print(len(d))
    # pitch_min = 100
    # pitch_max = 0
    # yaw_min = 100
    # yaw_max = 0
    # land_x_min = 100
    # land_x_max = 0
    # land_y_min = 100
    # land_y_max = 0
    # for i in range(len(d)):
    #     (data, label) = d.__getitem__(i)
    #     target_reg = label['target_reg']
    #     landmarks = label['landmarks']
    #     pitch = target_reg[0]
    #     yaw = target_reg[1]
    #     if pitch > pitch_max:
    #         pitch_max = pitch
    #     if pitch < pitch_min:
    #         pitch_min = pitch
    #     if yaw > yaw_max:
    #         yaw_max = yaw
    #     if yaw < yaw_min:
    #         yaw_min = yaw
    #     for j in range(landmarks.shape[0]):
    #         x = landmarks[j][0]
    #         y = landmarks[j][1]
    #         if x > land_x_max:
    #             land_x_max = x
    #         if x < land_x_min:
    #             land_x_min = x
    #         if y > land_y_max:
    #             land_y_max = y
    #         if y < land_y_min:
    #             land_y_min = y
    # print(f"pitch_min:{pitch_min} pitch_max:{pitch_max} yaw_min:{yaw_min} yaw_max:{yaw_max}")
    # print(f"landmarks_x_min:{land_x_min} landmarks_x_max:{land_x_max}model_old land_y_min:{land_y_min} land_y_max:{land_y_max}")
    # # pitch_min:-98.1966781616211 pitch_max:169.99476623535156 yaw_min:-79.2773208618164 yaw_max:23.500993728637695
    # # landmarks_x_min:16.600000381469727 landmarks_x_max:229.6999969482422 land_y_min:4.300000190734863 land_y_max:170.0